# --*-- coding: utf-8 --*--

from pika import BlockingConnection, ConnectionParameters, \
    PlainCredentials
from pika.exceptions import AMQPConnectionError
from functools import wraps
from time import sleep
from json import dumps

from pymongo import MongoClient, errors
from pymongo.collection import Collection
from pymongo.database import Database
from redis import Redis, exceptions
from retry.api import retry_call

from common import log

r = 0

# redis_retry_err = [exceptions.ConnectionError]

"""
数据库重连
:param func_name: 函数名
:param errs: 重连报错类
:param repeat_count:
:param t: 数字重连次数 false无限重连
:return: object
"""


def retry_connect(func_name, errs, repeat_count=None, t=5):
    global r
    repeat_count_type = type(repeat_count)
    r = repeat_count if repeat_count_type == int else True

    def wraps_func(func):
        @wraps(func)
        def method(*args, **kwargs):
            global r
            try:
                return func(*args, **kwargs)
            except errs as e:
                log.info(f"sleep {t}s, err:{e}")
                sleep(t)
                self = list(args)[0]
                connect_func = getattr(self, func_name)
                connect_func()

                if repeat_count_type == int and r > 0:
                    log.info(f"还有{r}次重试，报错:{e}")
                    r -= 1
                elif r <= 0:
                    raise Exception("达到最大重试次数")

                method(*args, **kwargs)

        return method

    return wraps_func


"""
redis无限读取队列模式，可与rabbitmq共用
:return: object
"""


def repeat_get_data():
    def wraps_func(func):
        @wraps(func)
        def method(*args, **kwargs):
            def _method():
                values = func(*args, **kwargs)
                if values is None:
                    is_repeat = kwargs.get("is_repeat", False)
                    args_len = len(args)
                    if args_len >= 3:
                        is_repeat = is_repeat if is_repeat else args[2]
                    if is_repeat:
                        raise AttributeError
                return values

            return retry_call(_method, delay=10, exceptions=AttributeError)

        return method

    return wraps_func


class RedisTools:

    def __init__(self, host, port, db, password):
        self.password = password
        self.host = host
        self.port = port
        self.db = db
        self.rds_cnt = None
        self.set_redis_cnt()

    def set_redis_cnt(self):
        """
        重连函数
        :return: object
        """
        if self.rds_cnt:
            self.rds_cnt.close()
        if self.password:
            self.rds_cnt = Redis(host=self.host, port=self.port, db=self.db, password=self.password)
        else:
            self.rds_cnt = Redis(host=self.host, port=self.port, db=self.db)

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    def push_set_data(self, keys, values):
        """
        获取set 队列数据
        :param keys:
        :param values:
        :return: object
        """
        self.rds_cnt.sadd(keys, values)

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    def push_list_data(self, keys, values):
        """
        获取list 队列数据
        :param keys:
        :param values:
        :return: object
        """
        self.rds_cnt.lpush(keys, values)

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    def is_member(self, keys, values):
        """
        keys队列是否存在values
        :param keys:
        :param values:
        :return: boole
        """
        return self.rds_cnt.sismember(keys, values)

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    @repeat_get_data()
    def get_set_data(self, keys, is_repeat=False):
        """
        获取set 队列数据
        :param keys:
        :param is_repeat:
        :return:
        """
        values = self.rds_cnt.spop(keys)
        return values.decode("utf8") if values else None

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    @repeat_get_data()
    def get_list_data(self, keys, is_repeat=False):
        """
        获取list队列数据
        :param keys:
        :param is_repeat:
        :return:
        """
        values = self.rds_cnt.lpop(keys)
        return values.decode("utf8") if values else None

    @retry_connect("set_redis_cnt", exceptions.ConnectionError)
    def get_lens(self, kye):
        """
        获取队列长度
        :param kye:
        :return:
        """
        return self.rds_cnt.scard(kye)

    def del_db(self, kye):
        """
        删除某个队列
        :param kye:
        :return:
        """
        return self.rds_cnt.delete(kye)

    def __del__(self):
        if self.rds_cnt:
            self.rds_cnt.close()


class MongoTools:

    def __init__(self, mongo_client, dbs="crawl_info"):
        self.mg_ct = None
        self.mg_db = None
        self.mg_tb = None
        self.host = mongo_client
        self.dbs = dbs
        self.mg_connect()

    def mg_connect(self):
        self.mg_ct = MongoClient(self.host)
        self.mg_db = Database(self.mg_ct, self.dbs)

    def table_connect(self, table):
        self.mg_tb = Collection(self.mg_db, table)

    # 统计条数
    @retry_connect('mg_connect', errors.ServerSelectionTimeoutError)
    def get_count(self, find=None, table="xiaomo_tools"):
        """
        统计条数
        :param find: 默认查全表条数，可以过滤
        :param table:
        :return:
        """
        find = find if find else dict()
        self.table_connect(table)
        return self.mg_tb.count_documents(find)

    # 获取数据
    @retry_connect('mg_connect', errors.ServerSelectionTimeoutError)
    def get_data(self, find=None, over=None, limit=0, skip=0, table="xiaomo_tools"):
        """
        获取数据
        :param find: 查询条件
        :param over: 显示字段
        :param limit:
        :param skip:
        :param table:
        :return:
        """
        find = find if find else dict()
        over = over if over else {"abcd_efg_hij_klm": 0}
        self.table_connect(table)
        if limit == 0:
            return list(self.mg_tb.find(find, over))
        else:
            return list(self.mg_tb.find(find, over).limit(limit).skip(skip))

    # 聚合获取数据
    @retry_connect('mg_connect', errors.ServerSelectionTimeoutError)
    def get_aggregate_data(self, find=None, table="xiaomo_tools"):
        """
        聚合函数
        :param find:
        :param table:
        :return:
        """
        find = find if find else list()
        self.table_connect(table)
        text = self.mg_tb.aggregate(find)
        return list(text)


def callback(channel, method, properties, body):
    print(body)
    channel.basic_ack(delivery_tag=method.delivery_tag)


class RabbitmqTool:
    def __init__(self, conf):
        self.host = conf.get('host', None)
        self.port = conf.get('port', None)
        self.user = conf.get('user', None)
        self.password = conf.get('password', None)
        self.exchange = conf.get('exchange', None)
        self.queue = conf.get('queue', None)
        self.virtual_host = conf.get('virtual_host', None)

        self.connect = None
        self.channel = None
        self.rbt_connect()

    def rbt_connect(self):
        """
        创建连接
        :return:
        """
        self.connect = BlockingConnection(
            ConnectionParameters(host=self.host, virtual_host=self.virtual_host, port=self.port,
                                 credentials=PlainCredentials(self.user, self.password)))
        self.channel = self.connect.channel()

    def create_queue(self, queue_name):
        """
        声明队列
        :param queue_name:
        :return:
        """
        self.channel.queue_declare(queue=queue_name, durable=True)

    def create_exchange(self, exchange_name):
        """
        声明交换器
        :param exchange_name:
        :return:
        """
        self.channel.exchange_declare(exchange=exchange_name)

    def queue_bind_exchange(self, queue=None, exchange=None):
        """
        队列绑定虚拟机
        :param queue:
        :param exchange:
        :return:
        """
        exchange = exchange if exchange else self.exchange
        queue = queue if queue else self.queue
        self.channel.queue_bind(exchange=exchange, queue=queue)

    @retry_connect('rbt_connect', AMQPConnectionError)
    def push_data(self, text, exchange=None, queue=None):
        """
        上传数据
        :param text:
        :param exchange:
        :param queue:
        :return:
        """
        exchange = exchange if exchange else self.exchange
        queue = queue if queue else self.queue
        log.info(f'{exchange}，{queue}推送资源：{text}')
        self.channel.basic_publish(exchange=exchange, routing_key=queue, body=dumps(text, ensure_ascii=False))

    @retry_connect('rbt_connect', AMQPConnectionError)
    def pull_data(self, queue=None, call_back=None):
        """
        获取数据
        :param queue:
        :param call_back:
        :return:
        """
        call_back = call_back if call_back else callback
        queue = queue if queue else self.queue
        self.channel.basic_consume(
            queue=queue,
            auto_ack=False,
            on_message_callback=call_back)
        self.channel.start_consuming()


if __name__ == '__main__':
    from common.base import data_base_conf

    conf = data_base_conf["connect"]
    MongoTools(conf)
